package program.leetcode.queue;

import java.util.*;

/**
 * 347. 前 K 个高频元素
 * <p>
 * 示例 1:
 * <p>
 * 输入: nums = [1,1,1,2,2,3], k = 2
 * 输出: [1,2]
 * <p>
 * 示例 2:
 * <p>
 * 输入: nums = [1], k = 1
 * 输出: [1]
 * <p>
 * 说明：
 * <p>
 * 你可以假设给定的 k 总是合理的，且 1 ≤ k ≤ 数组中不相同的元素的个数。
 * 你的算法的时间复杂度必须优于 O(n log n) , n 是数组的大小。
 * <p>
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode-cn.com/problems/top-k-frequent-elements
 * 著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。
 *
 *
 * 底层数据结构：PriorityQueue， java底层使用的是最小堆.
 * 直接存入元素.
 */
public class LeetCode_347_V2 {
    public List<Integer> topKFrequent(int[] nums, int k) {
        TreeMap<Integer, Integer> treeMap = new TreeMap<>();
        // 1.得到每个元素对应的频次.
        for (int num : nums) {
            if (treeMap.containsKey(num)) {
                treeMap.put(num, treeMap.get(num) + 1);
            } else {
                treeMap.put(num, 1);
            }
        }

        // 2. 使用优先队列, 得到k个频次最高的元素.
        PriorityQueue<Integer> pq = new PriorityQueue<>(Comparator.comparingInt(treeMap::get));
        for (Integer key : treeMap.keySet()) {
            if (pq.size() < k) {
                // 不满k个元素的时候, 直接放入优先队列中
                pq.add(key);
            } else if (treeMap.get(key) > treeMap.get(pq.peek())) {
                /**
                 * 当优先队列已有k个元素的时候;
                 * 获取第一个元素(也是频次最少的元素), 比较两者的频次;
                 * 若当前元素频次大于已存在优先队列中的最少频次的元素, 则出队最少频次元素; 入队次元素.
                 */
                pq.remove();
                pq.add(key);
            }
        }

        // 3.放入到List中, 返回数据
        List<Integer> res = new LinkedList<>();
        while (!pq.isEmpty()) {
            res.add(pq.remove());
        }

        return res;
    }

    public static void main(String[] args) {
        PriorityQueue<Integer> pq = new PriorityQueue<>();
        pq.add(1);
        pq.add(2);
        pq.add(3);
        pq.add(4);

        System.out.println(pq);
        System.out.println(pq.poll());
        System.out.println(pq);
        System.out.println(pq.peek());
        System.out.println(pq);
    }
}
